'''
https://leetcode.cn/problems/kth-largest-element-in-an-array
'''
import heapq
import random
from typing import List


class Solution:
    # 门槛堆(小根堆), O(n * log(K))
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = nums[:k]
        heapq.heapify(heap)
        for i in range(k, len(nums)):
            if nums[i] > heap[0]:
                heapq.heapreplace(heap, nums[i])    # pop and push
        return heap[0]

    # O(N)
    # 荷兰国旗，partation, 优化快排的那个partation
    # 只会往一侧走
    def findKthLargest2(self, nums: List[int], k: int) -> int:
        k -= 1
        def partation(l, r, x):
            # 倒序，[l, left_board)都比x大，(right_board ,r]右边都比x小，[left_board, right_board]都是x
            left_board, right_board = l, r
            i = l
            while i <= right_board:
                if nums[i] < x:
                    nums[i], nums[right_board] = nums[right_board], nums[i]
                    right_board -= 1    # i 不懂，这个元素继续看
                elif nums[i] > x:
                    nums[i], nums[left_board] = nums[left_board], nums[i]
                    i += 1
                    left_board += 1
                else:
                    i += 1      # 边界不动
            return left_board, right_board
        def process(l, r):
            x = nums[random.randint(l, r)]
            left_board, right_board = partation(l, r, x)
            if left_board <= k <= right_board:
                return nums[k]
            if k < left_board:
                return process(l, left_board-1)
            else:
                return process(right_board+1, r)
        return process(0, len(nums)-1)


nums = [3,2,3,1,2,4,5,5,6]
k = 4
print(Solution().findKthLargest2(nums, k))